package com.alameda.common.filter;

import lombok.extern.slf4j.Slf4j;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.slf4j.MDC;

import java.util.UUID;

import static com.alameda.common.base.constant.ProjectConstant.TRACE_ID;

/**
 * traceId过滤器 用于在Dubbo的RPC调用过程中传递和生成traceId，以实现请求链路追踪。
 *
 * @author Lv
 * @date 2024/7/8
 */
@SuppressWarnings("all")
@Slf4j
@Activate(group = {CommonConstants.PROVIDER, CommonConstants.CONSUMER})
public class DubboTraceIdFilter implements Filter {
    /**
     * 拦截处理rpc调度
     *
     * @param invoker    调用者信息，包含服务提供者地址、接口信息
     * @param invocation 调用信息，包含方法名、参数
     * @return
     */
    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) {
        RpcContext rpcContext = RpcContext.getContext();
        Result result = null;
        try {
            String traceId = rpcContext.getAttachment(TRACE_ID);
            if (traceId == null) {
                // 通过MDC获取上下文中的traceId
                if (MDC.get(TRACE_ID) != null) {
                    traceId = MDC.get(TRACE_ID);
                } else {
                    // 重新生成
                    traceId = UUID.randomUUID().toString().replace("-", "");
                }
            }

            // 设置附加参数
            rpcContext.setAttachment(TRACE_ID, traceId);
            // 设置mdc 用于日志打印
            MDC.put(TRACE_ID, traceId);
            result = invoker.invoke(invocation);
            if (rpcContext.isConsumerSide()) {
                log.info("consumer traceId:{}", traceId);
            }
            if (rpcContext.isProviderSide()) {
                log.info("provider traceId:{}", traceId);
            }
        } finally {
            if (rpcContext.isProviderSide()) {
                // provider端调用完毕收移除mdc值
                rpcContext.removeAttachment(TRACE_ID);
                if (log.isDebugEnabled()) {
                    log.debug("TraceIdFilter traceId and userId clear");
                }
            }
        }
        return result;
    }
}
